import functools, inspect, threading, time, uuid, traceback, zmq, json, re
try:
    from .data import qualify
    from .type_schema import TypeSchema
    from .task import SymbolicTask
except ImportError:
    from data import qualify
    from type_schema import TypeSchema
    from task import SymbolicTask
esc_sym = "\\|+-[](){}.*$^"
task_status = {}
def escaped(s):
    for c in esc_sym:
        s.replace(c, '\\'+c)
    return s
def edp_service_zmq(cls):
    def _update_progress(tid, percentage, status=None):
        pi = cls.pi if cls.pi is not None else cls.ppi
        msg = f"{percentage},{status}" if status is not None else f"{percentage},"
        pi.send_string(f"{tid}#" + msg)
        if status:
            global task_status
            task_status[tid]["sts"] = status
    def update_progress(self, percentage, status=None):
        if not hasattr(self, 'task_id'):
            assert 0 == 1, "call `update_progress()` from unlinked environment"
        return _update_progress(self.task_id, percentage, status)
    def _log(tid, msg: str):
        pi = cls.pi if cls.pi is not None else cls.lpi
        pi.send_string(f"#{tid}#" + msg)
    def log(self, message):
        if not hasattr(self, 'task_id'):
            assert 0 == 1, "call `log()` from unlinked environment"
        return _log(self.task_id, message)
    def name():
        if hasattr(cls, 'name') and isinstance(cls.name, str):
            return cls.name
        elif hasattr(cls, 'name') and callable(cls.name):
            return cls.name()
        else:
            return cls.__name__
    def desc():
        if hasattr(cls, 'description') and isinstance(cls.description, str):
            return cls.description
        elif hasattr(cls, 'description') and callable(cls.description):
            return cls.description()
        elif hasattr(cls, 'desc') and isinstance(cls.desc, str):
            return cls.desc
        elif hasattr(cls, 'desc') and callable(cls.desc):
            return cls.desc()
        else:
            return cls.__name__
    def category():
        if hasattr(cls, 'category') and isinstance(cls.category, str):
            return cls.category
        elif hasattr(cls, 'category') and callable(cls.category):
            return cls.category()
        else:
            return None
    def tasks_nominal():
        if hasattr(cls, 'tasks') and isinstance(cls.tasks, list):
            return cls.tasks
        elif hasattr(cls, 'tasks') and callable(cls.tasks):
            return cls.tasks()
        else:
            return []
    def tasks_symbolic():
        try:
            res = {}
            for task in tasks_nominal():
                sym = SymbolicTask.try_parse(task)
                res[sym.name] = sym
            return res
        except Exception as e:
            print(f"task signiture parse error: {e}")
            return {}
    tsblc = tasks_symbolic()
    def workflows():
        if hasattr(cls, 'workflows') and isinstance(cls.workflows, list):
            return cls.workflows
        elif hasattr(cls, 'workflows') and callable(cls.workflows):
            return cls.workflows()
        else:
            return []
    def dual_pipe(pp, lp):
        context = zmq.Context()
        pso = context.socket(zmq.PUB)
        pso.bind(f"tcp://*:{pp}")
        lso = context.socket(zmq.PUB)
        lso.bind(f"tcp://*:{lp}")
        print(f'bind progress socket at {pp} and log socket at {lp}')
        return pso, lso
    def single_pipe(pp):
        context = zmq.Context()
        pso = context.socket(zmq.PUB)
        pso.bind(f"tcp://*:{pp}")
        print(f'bind progress+log socket at {pp}')
        return pso
    @classmethod
    def run(cls, port, progress_port, log_port=None):
        context = zmq.Context()
        socket = context.socket(zmq.REP)
        socket.bind(f"tcp://*:{port}")
        if log_port: cls.pi = None; cls.ppi, cls.lpi = dual_pipe(progress_port, log_port)
        else: cls.pi = single_pipe(progress_port); cls.ppi = cls.lpi = None
        cls.log_port = log_port or progress_port
        cls.prog_port = progress_port
        print(f'bind main REP socket at {port}')
        while True:
            message = socket.recv_string() #[1:] # skip leading '$'
            rep = dispatch(message)
            socket.send_string(rep)
    cls.run = run
    cls.update_progress = update_progress
    cls.log = log
    tasks_names = tsblc.keys()
    is_instant = {}
    task_entries = {}
    def call_func(func, inputs, instance):
        if isinstance(func, staticmethod):
            res = func.__get__(None, cls)(inputs)
        elif isinstance(func, classmethod):
            res = func.__get__(None, cls)(cls, inputs)
        else:
            res = func(instance, inputs)
        return res
    for tn in tasks_names:
        if not hasattr(cls, tn):
            print(f'no task func: {tn}, skip route building')
            continue
        func = getattr(cls, tn)
        if hasattr(func, '_is_instant_task') and func._is_instant_task:
            is_instant[tn] = True
            def _start(inputs, tn):
                instance = cls()
                return call_func(getattr(cls, tn), inputs, instance)
            task_entries[tn] = _start
        else:
            def _start(inputs, tn):
                task_id = str(uuid.uuid4())
                global task_status
                task_status[task_id] = {
                    "sts": "Created",
                    "method": tn,
                    "created_at": time.time(),
                    "completed_at": None,
                    "res": None,
                    "err": None,
                }
                def run_task(task_id):
                    global task_status
                    _update_progress(task_id, 0.0, "Running")
                    instance = cls()
                    instance.task_id = task_id
                    ti = task_status[task_id]
                    try:
                        result = call_func(getattr(cls, tn), inputs, instance)
                        _update_progress(task_id, 100.0, "Completed")
                        ti["res"] = result
                        ti["completed_at"] = time.time()
                    except Exception as e:
                        stack_trace = traceback.format_exc()
                        print(f"{tn} task error:\n{stack_trace}")
                        error_msg = str(e)
                        _update_progress(task_id, None, "Failed")
                        ti["err"] = error_msg
                        ti["completed_at"] = time.time()
                thread = threading.Thread(target=run_task, args=[task_id])
                thread.daemon = True
                thread.start()
                return task_id
            is_instant[tn] = False
            task_entries[tn] = _start

    pat_task_result = re.compile(rf'/task/result/[a-z0-9-]+')
    def for_task_result(req):
        return pat_task_result.fullmatch(req) is not None
    def dispatch(req):
        print(f"received: {req}")
        if req == "/":
            return desc()
        elif req == "/meta":
            return json.dumps({
                "name": name(),
                "desc": desc(),
                "category": category()
            })
        elif req == "/tasks":
            return json.dumps(tasks_nominal())
        elif req == "/workflows":
            return json.dumps(workflows())
        elif for_task_result(req):
            return dispatch_task_result(req)
        return dispatch_start_task(req) or '#404'
    def dispatch_task_result(req):
        tid = req[13:]
        global task_status
        if tid not in task_status:
            return '#NoID'
        ti = task_status[tid]
        if ti["sts"] == "Completed":
            return ti["res"]
        elif ti["sts"] == "Failed":
            return '#TF#' + ti["err"]
        else:
            return '#TR'
    def dispatch_start_task(req):
        tn, inputs = req[6:].split('#', 1)
        if tn not in tsblc or tn not in is_instant or tn not in task_entries:
            return '#NoT'
        inputs = json.loads(inputs)
        if not qualify(inputs, tsblc[tn].input_type):
            return '#IQF'
        entry = task_entries[tn]
        if is_instant[tn]:
            try:
                res = entry(inputs, tn)
                return json.dumps({
                    "task_id": 0,
                    "progress_pipe_addr": "",
                    "result_addr": "",
                    "instant_result": res,
                    "status": "ResultReceived"
                })
            except Exception as e:
                stack_trace = traceback.format_exc()
                print(f"instant '{tn}' task error:\n{stack_trace}")
                return json.dumps({
                    "task_id": 0,
                    "progress_pipe_addr": "",
                    "result_addr": "",
                    "instant_result": None,
                    "status": "Failed",
                    "error_message": str(e)
                })
        else:
            task_id = entry(inputs, tn)
            return json.dumps({
                "task_id": task_id,
                "progress_pipe_addr": f"#{cls.prog_port}",
                "result_addr": f"/task/result/{task_id}",
                "log_pipe_addr": f"#{cls.log_port}",
            })
    return cls